import math
from data_loader import load_iris_data, train_test_split

class KNN:
    """K近邻算法实现"""
    
    def __init__(self, k=3):
        self.k = k
        self.X_train = None
        self.y_train = None
    
    def euclidean_distance(self, x1, x2):
        """计算欧几里得距离"""
        distance = 0.0
        for i in range(len(x1)):
            distance += (x1[i] - x2[i]) ** 2
        return math.sqrt(distance)
    
    def fit(self, X, y):
        """存储训练数据"""
        self.X_train = X
        self.y_train = y
    
    def predict(self, X):
        """预测类别"""
        predictions = []
        for x in X:
            prediction = self._predict_single(x)
            predictions.append(prediction)
        return predictions
    
    def _predict_single(self, x):
        """单个样本预测"""
        if self.X_train is None or self.y_train is None:
            raise ValueError("模型尚未训练")
        
        # 计算所有距离
        distances = []
        for i, x_train in enumerate(self.X_train):
            dist = self.euclidean_distance(x, x_train)
            distances.append((dist, i))
        
        # 按距离排序
        distances.sort(key=lambda x: x[0])
        
        # 获取最近的k个样本
        k_nearest = distances[:self.k]
        
        # 统计k个最近邻的标签
        class_votes = {}
        for dist, idx in k_nearest:
            label = self.y_train[idx]
            if label in class_votes:
                class_votes[label] += 1
            else:
                class_votes[label] = 1
        
        # 返回票数最多的标签
        max_votes = 0
        predicted_class = None
        for label, votes in class_votes.items():
            if votes > max_votes:
                max_votes = votes
                predicted_class = label
        
        return predicted_class
    
    def accuracy(self, y_true, y_pred):
        """计算准确率"""
        if len(y_true) != len(y_pred):
            raise ValueError("真实标签和预测标签数量不一致")
        
        correct = 0
        for i in range(len(y_true)):
            if y_true[i] == y_pred[i]:
                correct += 1
        
        return correct / len(y_true)

def main():
    print("=== KNN算法 - Iris分类（纯Python实现） ===\n")
    
    # 加载数据
    X, y, class_names = load_iris_data()
    print(f"数据记录数: {len(X)}")
    print(f"特征数量: {len(X[0])}")
    print(f"类别: {class_names}")
    print(f"类别分布: ", end="")
    for i, name in enumerate(class_names):
        count = y.count(i)
        print(f"{name}: {count}", end="  ")
    print()
    
    # 划分训练测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
    print(f"\n训练集大小: {len(X_train)}")
    print(f"测试集大小: {len(X_test)}")
    
    # 测试不同的k值
    k_values = [1, 3, 5, 7, 9]
    best_accuracy = 0
    best_k = 0
    
    print("\n不同k值的准确率:")
    print("k值\t准确率")
    print("-" * 20)
    
    for k in k_values:
        knn = KNN(k=k)
        knn.fit(X_train, y_train)
        y_pred = knn.predict(X_test)
        accuracy = knn.accuracy(y_test, y_pred)
        
        print(f"{k}\t{accuracy:.4f}")
        
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            best_k = k
    
    print(f"\n最佳k值: {best_k}, 最高准确率: {best_accuracy:.4f}")
    
    # 使用最佳k值显示详细结果
    print(f"\n使用最佳k值({best_k})的详细预测结果:")
    print("样本\t真实类别\t预测类别\t是否正确")
    print("-" * 45)
    
    knn_best = KNN(k=best_k)
    knn_best.fit(X_train, y_train)
    y_pred_best = knn_best.predict(X_test)
    
    correct_count = 0
    for i in range(min(15, len(y_test))):
        true_label = class_names[y_test[i]]
        pred_label = class_names[y_pred_best[i]]
        correct = "✓" if y_test[i] == y_pred_best[i] else "✗"
        if correct == "✓":
            correct_count += 1
        
        print(f"{i+1}\t{true_label}\t{pred_label}\t{correct}")
    
    print(f"\n前{min(15, len(y_test))}个样本中正确分类: {correct_count}个")
    
    # 计算混淆矩阵
    print("\n混淆矩阵:")
    n_classes = len(class_names)
    confusion_matrix = [[0] * n_classes for _ in range(n_classes)]
    
    for i in range(len(y_test)):
        true_idx = y_test[i]
        pred_idx = y_pred_best[i]
        confusion_matrix[true_idx][pred_idx] += 1
    
    # 打印混淆矩阵
    print("    " + " ".join(f"{name:>8}" for name in class_names))
    for i, name in enumerate(class_names):
        print(f"{name:4}" + " ".join(f"{confusion_matrix[i][j]:8}" for j in range(n_classes)))

if __name__ == "__main__":
    main()